import numpy as np

from calc_force import *
from read_data import *
from write_data import *
from plot_data import *
from filter_data import *
from transform_coord import *
from check_exist import *
from time_cut import *
from globals import *

# 0: 不保存，1: 保存, 2: 数据集
save_flag = 2

# F16/F22/F22_STATIC/F22_10/F22_SWEEP/FlyWing/
MODEL = "F22"  

DATASET_PREFIX = "../DATASET_ALL/DATASET_UNCUT/"
AERO_DIR = DATASET_PREFIX + MODEL + "/AERO/"
ANGLE_DIR = DATASET_PREFIX + MODEL + "/ANGLE/"
TENSION_DIR = DATASET_PREFIX + MODEL + "/TENSION/"
ENCODER_DIR = DATASET_PREFIX + MODEL + "/ENCODER/"
OUTPUT_DIR = "../OUTPUT_Filter/"
OUTPUT_DATASET_DIR = "../OUTPUT_Dataset/"

# 读取文件列表
print("=========项目读取=========")
# file_list = get_l_folder(INPUT_DIR)
file_list = get_l_list(MODEL) 

# 获取频率信息
if(MODEL=="F22" or MODEL=="F22_10"):
    freq_motion_dict = get_f(MODEL, DATASET_PREFIX)
freq_aero = 100
freq_tension = 10
freq_encoder = 20
freq_angle = 200

# file_list = ["A04"]

# 气动力计算迭代轮数
iter_num = 5


# 检查文件完整性
print("========完整性检查========")
file_list = check_e(file_list, MODEL, DATASET_PREFIX)


for file_name in file_list:
    print("==========="+file_name+"===========")

    # 读取气动数据（电压值）
    u_nw_raw = read_d_aero(AERO_DIR + file_name + "_NW.txt")
    u_w_raw = read_d_aero(AERO_DIR + file_name + "_W.txt")

    # 读取拉力数据
    tension_nw = read_d_tension(TENSION_DIR + file_name + "_NW*.txt")
    tension_w = read_d_tension(TENSION_DIR + file_name + "_W*.txt")


    # 读取编码器数据
    encoder_nw = read_d_encoder(ENCODER_DIR + file_name + "_NW*.txt")
    encoder_w = read_d_encoder(ENCODER_DIR + file_name + "_W*.txt")

    # 读取姿态角数据
    angle_nw = read_d_angle(ANGLE_DIR + file_name + "_NW*.csv")
    angle_w = read_d_angle(ANGLE_DIR + file_name + "_W*.csv")

    angle_nw = filter_angle(angle_nw)
    angle_w = filter_angle(angle_w)

    # 时间对齐
    freq_angle = 200
    u_nw_raw, tension_nw, encoder_nw, angle_nw = time_cut(u_nw_raw, freq_aero, tension_nw, freq_tension, encoder_nw, freq_encoder, angle_nw, freq_angle)
    u_w_raw, tension_w, encoder_w, angle_w = time_cut(u_w_raw, freq_aero, tension_w, freq_tension, encoder_w, freq_encoder, angle_w, freq_angle)
    freq_angle = freq_aero    # 经过裁剪
    u_nw_raw,u_w_raw,angle_w,angle_nw = time_cut2(u_nw_raw,u_w_raw,angle_w,angle_nw)

    
    # 气动力（电压）滤波
    if(MODEL=="F22" or MODEL=="F22_10"):
        freq_motion = freq_motion_dict[file_name]
    else:
        freq_motion = [0.0001,0,0,0]
    print("运动频率: "+str(freq_motion)+"Hz")

    u_nw = filter_aero(u_nw_raw, freq_aero, freq_motion)
    u_w = filter_aero(u_w_raw, freq_aero, freq_motion)
    # 不滤波
    # u_nw = u_nw_raw
    # u_w = u_w_raw

    # 角度滤波
    angle_nw = filter_aero(angle_nw, freq_aero, freq_motion)
    angle_w = filter_aero(angle_w, freq_aero, freq_motion)
    
    F_n_Q_aero = []
    F_n_Q_body = []
    F_n_Q_wind = []
    F_n_Q_rslt = []

    # 计算气动力(机体轴)
    for i in range(len(u_nw)):
        
        u0 = u_nw[i]
        u1 = u_w[i]

        # X, Y, Z, Mx, My, Mz, Q
        f,q = calc_f2(u0, u1, iter_num)

        f = np.append(f,q)
        F_n_Q_aero.append(f)

    F_n_Q_aero = np.array(F_n_Q_aero)

    # 坐标系转换(天平轴→体轴系)
    F_n_Q_body = transform_a2b(F_n_Q_aero)
    
    # 横侧向归零
    F_n_Q_body = filter_froce(F_n_Q_body)

    # 坐标系转换(体轴系→气流系)
    Angle_body = angle_w
    F_n_Q_wind, Angle_wind = transform_b2w(F_n_Q_body,Angle_body)
    F_n_Q_rslt = transform_w2r(F_n_Q_wind)

    # 横侧向归零
    F_n_Q_rslt = filter_froce(F_n_Q_rslt)

    # 计算气动力系数
    C_n_Q_body = calc_coeff(F_n_Q_body, MODEL)
    C_n_Q_wind = calc_coeff(F_n_Q_wind, MODEL)
    C_n_Q_rslt = calc_coeff(F_n_Q_rslt, MODEL)

    if(save_flag==0):
        # 画图
        plot_all(angle_nw,angle_w,freq_angle,C_n_Q_wind,C_n_Q_body,freq_aero,freq_motion,angle_w)
        # plot_all_time(angle_nw,angle_w,freq_angle,F_n_Q_body,F_n_Q_wind,freq_aero)
        # plot_all_time(angle_nw,angle_w,freq_angle,F_n_Q_aero,F_n_Q_aero,freq_aero)


        # 画三维图
        # plot_3d_trajectory(angle_nw,angle_w)
        # plot_3d_wind(F_n_Q_rslt, Angle_wind)

        # 时间-电压
        # plot_u(u_nw_raw, u_w_raw, u_nw, u_w, freq_aero)

        # 时间-角度
        # plot_comp(angle_nw,angle_w,freq_angle)

        # 时间-气动力
        # plot_f_time(C_n_Q_body, freq_aero)
        # plot_f_time(C_n_Q_wind, freq_aero)

        # 角度-气动力
        # plot_all_angle(angle_nw,angle_w,freq_angle,C_n_Q_rslt,C_n_Q_rslt,freq_aero,freq_motion,angle_w)
        # plot_f_degree_ani(F_n_Q_body, angle_w, freq_motion, freq_angle, False)
        # plot_f_degree(F_n_Q_aero, angle_w, freq_motion, freq_angle, False)

        # 气动角-气动力
        # plot_f_degree_wind(F_n_Q_wind, Angle_wind, freq_motion, freq_angle, True)
        # plot_f_degree_wind_all(C_n_Q_rslt, Angle_wind, freq_motion, freq_angle, False)

        # plot_all_angle_ani(angle_nw,angle_w,freq_angle,F_n_Q_body,F_n_Q_wind,freq_aero,freq_motion,angle_w)

    elif(save_flag==1):
        # 保存数据到文件
        force_list = [F_n_Q_aero, F_n_Q_body, F_n_Q_wind, F_n_Q_rslt]
        coeff_list = [C_n_Q_body, C_n_Q_wind, C_n_Q_rslt]
        anlge_list = [Angle_body, Angle_wind]
        pic_list = []
        all_list = [force_list, coeff_list, anlge_list, pic_list]
        write_all(all_list, MODEL,OUTPUT_DIR,file_name)
    
    elif(save_flag==2):
        # 创建数据集
        dt1 = 500
        dt2 = 200
        C_n_Q_body = C_n_Q_body[dt1:-dt2,]
        Angle_body = Angle_body[dt1:-dt2,]
        Angle_wind = Angle_wind[dt1:-dt2,]


        all_list = [C_n_Q_body, Angle_body, Angle_wind]
        write_dataset(all_list,MODEL,OUTPUT_DATASET_DIR,file_name)
        
        
    
    

